import os
import pickle
import numpy as np
import inspect
from PIL import Image, ImageOps

try:
	import gradio as gr
except Exception:
	gr = None


# 加载保存的 KNN 模型（保存为 (scaler, knn) 的元组）
def load_model(pickle_path=None):
	if pickle_path is None:
		pickle_path = os.path.join(os.path.dirname(__file__), 'best_knn_model.pkl')
	if not os.path.exists(pickle_path):
		raise FileNotFoundError(f"Model file not found: {pickle_path}")
	with open(pickle_path, 'rb') as f:
		obj = pickle.load(f)
	if not (isinstance(obj, tuple) and len(obj) == 2):
		raise ValueError('Pickle content is not (scaler, knn) tuple')
	return obj


# 预处理函数：将输入图像（PIL 或 numpy）转换为与 sklearn digits 数据一致的形状
def preprocess_image(img):
	# Debug: 输出原始输入类型和内容
	print(f"Debug: 原始输入类型: {type(img)}")
	if isinstance(img, dict):
		print(f"Debug: 原始输入内容 (dict keys): {list(img.keys())}")
		# 优先使用 composite 键
		if 'composite' in img:
			img = img['composite']
		else:
			raise ValueError("未找到 composite 键，无法处理输入")

	if isinstance(img, np.ndarray):
		print(f"Debug: 原始输入内容 (array shape): {img.shape}")
		# 处理 RGBA 图像
		if img.shape[2] == 4:  # 如果是 RGBA
			rgb, alpha = img[..., :3], img[..., 3]
			# 创建全白背景
			bg = np.ones_like(rgb) * 255
			# 仅在 alpha > 0 的地方保留原始 RGB
			mask = (alpha > 0)[..., None]
			img = np.where(mask, rgb, bg).astype('uint8')
		img = Image.fromarray(img)

	if not isinstance(img, Image.Image):
		raise ValueError(f"无法识别的图像输入类型: {type(img)}")

	# 转为灰度
	img = img.convert('L')
	# 将图像缩放到 8x8 的像素（sklearn digits 是 8x8）
	img = img.resize((8, 8), Image.BILINEAR)

	# 转换为 numpy，并归一化到类似 digits 数据的范围
	arr = np.asarray(img).astype(float)
	# 如果全白，提示用户写字
	if arr.max() < 10:
		raise ValueError('画板内容为空，请在画板上写字')
	# 自动判断是否反色：如果均值大于8，说明是白底黑字，需反色
	if arr.mean() > 8:
		arr = 255 - arr

	# Debug: 输出处理后的灰度图像统计信息
	print(f"Debug: 灰度图像统计信息 - min: {arr.min()}, max: {arr.max()}, mean: {arr.mean()}")
	print(f"Debug: 缩放后的8x8矩阵:\n{np.round(arr, 2)}")

	# sklearn digits 的像素范围大约在 0-16（整数），我们把 0-255 映射到 0-16
	arr = (arr / 255.0) * 16.0
	# 扁平化为 1x64
	arr = arr.reshape(1, -1)
	return arr


def make_predict_fn(model_tuple):
	scaler, knn = model_tuple

	def predict_from_image(img, debug=False):
		# img 来自 gradio 的 sketchpad 或 image upload
		try:
			x = preprocess_image(img)
			# 如果预处理返回非期望形状，主动报错
			if x.ndim != 2 or x.shape[1] != 64:
				return (None, f'预处理后形状异常: {x.shape}') if debug else f'Error: preprocessed shape {x.shape}'
			x_scaled = scaler.transform(x)
			pred = knn.predict(x_scaled)
			label = int(pred[0])
			if debug:
				stats = {
					'min': float(x.min()),
					'max': float(x.max()),
					'mean': float(x.mean()),
					'raw': np.round(x.reshape(8,8), 2).tolist()
				}
				return label, f"stats: min={stats['min']:.3f}, max={stats['max']:.3f}, mean={stats['mean']:.3f},\nraw8x8={stats['raw']}"
			return label
		except Exception as e:
			if debug:
				return None, f'Error: {e}'
			return f'Error: {e}'

	return predict_from_image


def build_gradio_app(model_path=None):
	if gr is None:
		raise RuntimeError('Gradio 未安装，请先: pip install gradio')

	model_tuple = load_model(model_path)
	predict_fn = make_predict_fn(model_tuple)

	# 支持 sketchpad (绘图) 与 图片上传两种输入，兼容不同 gradio 版本
	with gr.Blocks() as demo:
		gr.Markdown('# KNN 手写数字识别')
		with gr.Row():
			with gr.Column():
				sketch = None
				# 尝试创建 Sketchpad，部分 gradio 版本不接受 shape 参数
				if hasattr(gr, 'Sketchpad'):
					try:
						sig = inspect.signature(gr.Sketchpad)
						if 'shape' in sig.parameters:
							sketch = gr.Sketchpad(label='在此画数字 (或上传图片)', shape=(280, 280))
						else:
							sketch = gr.Sketchpad(label='在此画数字 (或上传图片)')
					except Exception:
						sketch = None

				# 上传图片组件作为备选
				upload = gr.Image(label='或上传图片', type='pil')
				btn = gr.Button('预测')
			with gr.Column():
				out = gr.Label(label='预测结果')
				debug_out = gr.Textbox(label='Debug 输出（特征统计）')
				debug_chk = gr.Checkbox(label='开启 Debug 输出', value=False)

		# 通用回调，接受可变数量的输入（兼容只有上传或同时有 sketch 的情况）
		def on_click_predict(*args):
			# 最后一个参数可能是 debug checkbox（bool），所以从右侧取值
			debug = False
			if len(args) > 0 and isinstance(args[-1], bool):
				debug = args[-1]
				inputs_imgs = args[:-1]
			else:
				inputs_imgs = args

			img = None
			for a in inputs_imgs:
				if a is not None:
					img = a
			if img is None:
				if debug:
					return None, '未提供图像'
				return '未提供图像', ''

			if debug:
				label, info = predict_fn(img, debug=True)
				return label, info
			else:
				result = predict_fn(img, debug=False)
				# 如果只返回一个值，补充空字符串作为第二输出
				if isinstance(result, tuple) and len(result) == 2:
					return result
				return result, ''

		# 将 debug checkbox 也作为输入（如果存在），并把 debug 输出放在第二个输出
		inputs = [c for c in (sketch, upload) if c is not None]
		if 'debug_chk' in locals():
			inputs.append(debug_chk)
			btn.click(on_click_predict, inputs=inputs, outputs=[out, debug_out])
		else:
			btn.click(on_click_predict, inputs=inputs, outputs=[out])

	return demo


if __name__ == '__main__':
	# 当作为脚本运行时，启动 gradio 服务
	try:
		if gr is None:
			print('请先安装 gradio: python -m pip install gradio')
		else:
			app = build_gradio_app()
			app.launch(share=False)
	except Exception as e:
		print('启动 webapp 时出错:', e)
